#include <vector>
#include <queue>
#include <iostream>
#include <xtensor.hpp>
#include <xtensor/xio.hpp>
#include <pybind11/pybind11.h>
#define FORCE_IMPORT_ARRAY
#include <xtensor-python/pyarray.hpp>

template <typename scalar_t>
xt::xtensor<npy_int32, 1> nms_cpu(xt::xtensor<scalar_t, 2>& boxes, xt::xtensor<scalar_t, 1>& scores, scalar_t threshold)
{
    auto x1 = xt::view(boxes, xt::all(), 0);
    auto y1 = xt::view(boxes, xt::all(), 1);
    auto x2 = xt::view(boxes, xt::all(), 2);
    auto y2 = xt::view(boxes, xt::all(), 3);
    auto areas = (x2 - x1 + 1) * (y2 - y1 + 1);
    auto order = xt::view(xt::argsort(scores), xt::range(_, _, -1));
    std::vector<size_t> reserved_box_idx;
    std::vector<size_t> order_vec{order.begin(), order.end()};
    while (!order_vec.empty()) {
        size_t i = order_vec[0];
        reserved_box_idx.push_back(i);
        auto keep_index = xt::view(xt::adapt(order_vec), xt::range(1, _));
        auto max_x1 = xt::maximum(x1(i), xt::view(x1, xt::keep(keep_index)));
        auto max_y1 = xt::maximum(y1(i), xt::view(y1, xt::keep(keep_index)));
        auto min_x2 = xt::minimum(x2(i), xt::view(x2, xt::keep(keep_index)));
        auto min_y2 = xt::minimum(y2(i), xt::view(y2, xt::keep(keep_index)));

        auto intersection = xt::maximum(0.0, min_x2 - max_x1 + 1) * xt::maximum(0.0, min_y2 - max_y1 + 1);
        auto areas_tail = xt::view(areas, xt::keep(keep_index));
        auto iou = intersection / (areas(i) + areas_tail - intersection + 1e-6);
        auto indexes = xt::where(iou <= threshold);
        std::vector<size_t>& selected = indexes[0];
        std::for_each(selected.begin(), selected.end(), [](size_t &i) { ++i; });
        // filter elements
        int j = 0;
        for (auto &idx : selected) {
            order_vec[j++] = order_vec[idx];
        }
        order_vec.erase(order_vec.begin() + j, order_vec.end());
    }
    return xt::adapt(reserved_box_idx);
}


PYBIND11_MODULE(cpp_numpy, m)
{
    xt::import_numpy();
    m.doc() = "Test module for xtensor python bindings";
    m.def("nms_cpu", nms_cpu<npy_float16>, "Non-maximum suppression, cpu version.");
    m.def("nms_cpu", nms_cpu<npy_float32>, "Non-maximum suppression, cpu version.");
    m.def("nms_cpu", nms_cpu<npy_float64>, "Non-maximum suppression, cpu version.");
}